from typing import Protocol, runtime_checkable
from fpowerkit import Grid
from .base import *

FILE_GEN = "gen"
FILE_BUS = "bus"
GEN_ATTRIB = ["P","Q","costp"]
GEN_TOT_ATTRIB = ["totP","totQ","totC"]
BUS_ATTRIB = ["Pd","Qd","Pg","Qg","cshd"]
BUS_TOT_ATTRIB = ["totPd","totQd","totPg","totQg"]

@runtime_checkable
class PluginWithGrid(Protocol):
    @property
    def Grid(self)->Grid: ...

def _find_grid_plugin(plugins:dict[str,PluginBase])->PluginWithGrid:
    for plg in plugins.values():
        if isinstance(plg, PluginWithGrid):
            return plg
    raise ValueError("未找到可以导出电网数据的插件")

class StaGen(StaBase):
    def __init__(self,path:str,suffix:str,tinst:TrafficInst,plugins:dict[str,PluginBase]):
        self.__plg = _find_grid_plugin(plugins)
        gen_names = self.__plg.Grid.GenNames
        super().__init__(FILE_GEN,path,suffix,cross_list(gen_names,GEN_ATTRIB)+GEN_TOT_ATTRIB,tinst,plugins)

    def GetData(self,inst:TrafficInst,plugins:list[PluginBase])->list:
        mpdn = self.__plg
        sb_MVA = mpdn.Grid.Sb_MVA
        _t = inst.current_time
        p = []; q = []; cp = []
        for g in mpdn.Grid.Gens:
            costthis = g.Cost(_t)
            if costthis is None:
                costthis = 0
            if g.P is None or g.Q is None:
                p.append(0); q.append(0)
                cp.append(costthis)
            else:
                p.append(g.P*sb_MVA); q.append(g.Q*sb_MVA)
                cp.append(costthis)
        return p + q + cp + [sum(p), sum(q), sum(cp)]

class StaBus(StaBase):
    def __init__(self,path:str,suffix:str,tinst:TrafficInst,plugins:dict[str,PluginBase]):
        self.__plg = _find_grid_plugin(plugins)
        bus_names = self.__plg.Grid.BusNames
        self.__bus_with_gens = [b.ID for b in self.__plg.Grid.Buses if len(self.__plg.Grid.GensAtBus(b.ID))>0]
        super().__init__(FILE_BUS,path,suffix,cross_list(bus_names,["Pd","Qd"]) 
            + cross_list(self.__bus_with_gens,["Pg","Qg","cshd"]) + BUS_TOT_ATTRIB,tinst,plugins)

    def GetData(self,inst:TrafficInst,plugins:dict[str,PluginBase])->list:
        mpdn = self.__plg.Grid
        sb_MVA = mpdn.Sb
        sb_kVA = mpdn.Sb*1000
        _t = inst.current_time
        Pd = [b.Pd(_t)*sb_MVA for b in mpdn.Buses]
        Qd = [b.Qd(_t)*sb_MVA for b in mpdn.Buses]
        Pg = []; Qg = []
        Cshadow = []
        for bn in self.__bus_with_gens:
            pg = 0; qg = 0
            for g in mpdn.GensAtBus(bn):
                if g.P is not None: pg += g.P
                if g.Q is not None: qg += g.Q
            Pg.append(pg*sb_MVA); Qg.append(qg*sb_MVA)
            sp = mpdn.Bus(bn).ShadowPrice
            Cshadow.append(sp/sb_kVA if sp is not None else 0)
        return Pd + Qd + Pg + Qg + Cshadow + [sum(Pd), sum(Qd), sum(Pg), sum(Qg)] # Unit = MVA